#!/usr/bin/python3
# -*- coding: UTF-8 -*-
from elasticsearch import Elasticsearch
from openpyxl import Workbook
import os
import datetime

url = "http://localhost:9200/"
index_name = "tmsearch"
doc_type = "tradeMark"
es = Elasticsearch(url, timeout=120)


# 日期格式化函数
def convert_date(date_str):
    if date_str is None:
        return ''
    date = datetime.datetime.strptime(date_str.replace('T', ' ')[0:19], "%Y-%m-%d %H:%M:%S")
    return (date + datetime.timedelta(hours=8)).strftime("%Y-%m-%d")


def convert_serve_list(serve_list):
    if serve_list is None:
        return ''
    list = []
    if not serve_list is None:
        for i in serve_list:
            if 'tmSimilarGroup' in i and 'tmGoodsCnName' in i:
                list.append(i['tmSimilarGroup'] + '-' + i['tmGoodsCnName'] + ';')
    if len(list) == 0:
        return ''
    else:
        return ''.join(list)


def query_trademark(_category):
    data = {
        "_source": ["tmAppNo", "tmCategoryNo", "tmAppDate", "tmName",
                    "tmRegisterNoticeDate", "tmRegistrant", "serveList"],
        "query": {
            "match_phrase": {
                "tmRegistrant.tmRegistrantCnName": "合肥龙图腾信息技术有限公司"
            }
        }
    }
    # 指定search_type="scan"模式，并返回_scroll_id给es.scroll获取数据使用
    query = es.search(index=index_name, doc_type=doc_type, body=data, scroll='5m', size=1000)
    results = query['hits']['hits']  # es查询出的结果第一页
    total = query['hits']['total']  # es查询出的结果总量
    scroll_id = query['_scroll_id']  # 游标用于输出es查询出的所有结果
    print(total)
    for i in range(0, int(total / 100) + 1):
        # scroll参数必须指定否则会报错
        query_scroll = es.scroll(scroll_id=scroll_id, scroll='5m')['hits']['hits']
        results += query_scroll

    for item in results:
        yield item['_source']['tmAppNo'], \
              item['_source']['tmCategoryNo'], \
              convert_date(item['_source']['tmAppDate']), \
              item['_source']['tmRegistrant']['tmRegistrantCnName'], \
              item['_source']['tmRegistrant']['tmRegistrantCnAddress'], \
              item['_source']['tmName'], \
              item['_source'].get('tmType', ''), \
              item['_source'].get('tmRegisterNoticeIssue', ''), \
              convert_date(item['_source'].get('tmRegisterNoticeDate', None)), \
              item['_source'].get('tmFirstTrialNoticeIssue', ''), \
              convert_date(item['_source'].get('tmFirstTrialNoticeDate', None)), \
              convert_date(item['_source'].get('tmSpecialIssueBeginDate', None)), \
              convert_date(item['_source'].get('tmSpecialIssueEndDate', None)), \
              item['_source'].get('tmSpecialValidity', ''), \
              '' if (item['_source'].get('tmProxyPerson') is None) else item['_source'].get('tmProxyPerson').get('tmProxyPersonCode', ''), \
              '' if (item['_source'].get('tmProxyPerson') is None) else item['_source'].get('tmProxyPerson').get('tmProxyPersonName', ''), \
              convert_serve_list(item['_source'].get('serveList', None)),



def generate_excel():
    wb = Workbook(write_only=True)
    ws = wb.create_sheet()
    ws.title = '合肥龙图腾信息技术有限公司商标信息'
    file_name = '合肥龙图腾信息技术有限公司商标信息.xlsx'
    result_file = os.path.join('E:/', file_name)
    if os.path.exists(result_file):
        os.remove(result_file)
    ws.append(['申请号', '分类号', '申请日期', '申请人', '申请人地址', '商标名称', '商标类型',
               '注册公告期号', '注册公告日期', '初审公告期号', '初审公告日期',
               '专用期开始日期', '专用期结束日期', '专用有效期', "代理人编码", "代理人名称", "服务列表"])
    category_list = [5]
    for category in category_list:
        for record in query_trademark(str(category)):
            print(record)
            ws.append(record)
    wb.save(filename=result_file)


if __name__ == "__main__":
    generate_excel()

